package main

import (
	"b_plus_tree/tree"
	"b_plus_tree/util"
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"
)

//挑战3：编码实现B+树存储结构，要支持插入、查询操作，
//实现后将5000万行KV数据（K为整数，V为1-100长度的字符串）导入其中，
//随机抽1000个K，在B+树中进行查询V，每条记录查询100次，求平均值，
//将所有查询结果写入文件，并在程序中输出平均查询时间。PK点：导入时间最短，查询时间最短，内存消耗最低。
//截止时间：2022年3月14日上午。
//注意：只能自己通过查询相关资料的方式独立完成，编码，在PK前不能请教其他人，不能再GitHub复制相关代码。
const (
	queryTime    = 100    //每个key查询次数
	randomKeyNum = 1000   //随机生成多少个key
	dataCount    = 500000 //总数据量大小，用于生成随机key
)

var (
	defaultDataFilePath   = "/data.txt"
	defaultKeyFilePath    = "/key.txt"
	defaultResultFilePath = "/result.txt"
	loadTime              time.Duration
	randomKeyOff          = 0
	randomKeys            = make(map[int]byte, randomKeyNum)
	result                = make(map[int][]string, randomKeyNum) //key -[value query time]
	dataPath              = flag.String("data", "", "write data file path ")
	resultPath            = flag.String("result", "", "write result file path")
	keyPath               = flag.String("key", "", "write key file path")
	memprofile            = flag.String("heap", "", "write memory profile to `file`")
	//树的阶
	m = flag.Int("m", 5, "steps of b plus tree")
)

func main() {
	//默认读取当前文件夹下的key和data
	getwd, err := os.Getwd()
	defaultDataFilePath = getwd + defaultDataFilePath
	defaultKeyFilePath = getwd + defaultKeyFilePath
	defaultResultFilePath = getwd + defaultResultFilePath
	flag.Parse()
	if err != nil {
		util.Log(err)
	}
	if *m != 0 {
		tree.M = *m
	}
	if *resultPath != "" {
		defaultResultFilePath = *resultPath
	}
	if *dataPath != "" {
		defaultDataFilePath = *dataPath
	}
	if *keyPath != "" {
		defaultKeyFilePath = *keyPath
	}
	_ = os.Remove(defaultResultFilePath)
	work()
}

func work() {
	myTree := new(tree.BPlusTree)
	file, err := os.OpenFile(defaultDataFilePath, os.O_RDWR, 0644)
	defer file.Close()
	if err != nil {
		util.Log(err)
	}
	reader := bufio.NewReader(file)
	myTree.Initialize()
	var insert int
	loadStart := time.Now()
	for {
		kv, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				if kv != "" {
					k, v := convertKv(kv)
					result := myTree.Insert(tree.Key(k), tree.Value(v))
					if result != false {
						insert++
					} else {
						fmt.Print(k, "-导入失败")
					}
				}
				break
			}
		}
		k, v := convertKv(kv)
		result := myTree.Insert(tree.Key(k), tree.Value(v))
		if result != false {
			insert++
		} else {
			fmt.Print(k, "-导入失败")
		}
	}
	since := time.Since(loadStart)
	loadResult := "导入时间:" + since.String() + " ," + "导入条数:" + strconv.Itoa(insert) + "\n"
	//导入结束，开始查询
	//随机生成1000个key
	queryAllTime := time.Now()
	keyFile, err := os.Open(defaultKeyFilePath)
	if err != nil {
		util.Log(err)
	}
	newReader := bufio.NewReader(keyFile)
	for {
		key, err := newReader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				if key != "" {
					key = strings.TrimSpace(key)
					atoi, err := strconv.Atoi(key)
					if err != nil {
						util.Log(err)
					}
					randomKeys[atoi] = 0
				}
				break
			} else {
				util.Log(err)
			}
		}
		key = strings.TrimSpace(key)
		atoi, err := strconv.Atoi(key)
		if err != nil {
			util.Log(err)
		}
		randomKeys[atoi] = 0
	}

	writeCount := 0
	for key := range randomKeys {
		writeCount++
		now := time.Now()
		write := true
		var value string
		for i := 0; i < queryTime; i++ {
			data, succ := myTree.FindData(tree.Key(key))
			if !succ {
				write = false
				break
			}
			if i == 0 {
				value = string(data)
			}
		}
		if write {
			avg := time.Since(now) / queryTime
			result[key] = append(result[key], []string{value, avg.String()}...)
		} else {
			result[key] = append(result[key], []string{value, "查询失败"}...)
		}
	}
	fmt.Println(writeCount)
	queryAllTimeSience := time.Since(queryAllTime)
	//写入结果
	resultFile, err := os.OpenFile(defaultResultFilePath, os.O_RDWR|os.O_CREATE, 0644)
	if err != nil {
		util.Log(err)
	}
	if err != nil {
		util.Log(err)
	}
	_, err = resultFile.WriteString(loadResult)
	head := "|line|--------|key|--------|value|--------|timeAvg(100)|\n"
	allQuryResult := "总查询耗时" + queryAllTimeSience.String() + "\n"
	_, err = resultFile.WriteString(time.Now().String() + "\n")
	_, err = resultFile.WriteString(allQuryResult)
	_, err = resultFile.WriteString(head)
	if err != nil {
		util.Log(err)
	}
	index := 1
	for key, valueAndSince := range result {
		line := strconv.Itoa(index) + ":" + strconv.Itoa(key) + " ," + valueAndSince[0] + ", " + valueAndSince[1] + "\n"
		index++
		_, err = resultFile.WriteString(line)
		if err != nil {
			util.Log(err)
		}
	}
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		defer f.Close()
		if err != nil {
			log.Fatal("could not create memory profile: ", err)
		}
		//runtime.GC() // get up-to-date statistics
		if err := pprof.Lookup("heap").WriteTo(f, 0); err != nil {
			log.Fatal("could not write memory profile: ", err)
		}
	}
}

func convertKv(kv string) (int, string) {
	input := strings.TrimSpace(kv)
	split := strings.Split(input, ",")
	if len(split) != 2 {
		util.Log(errors.New("数据有误"))
	}
	keyStr := strings.TrimSpace(split[0])
	key, err := strconv.Atoi(keyStr)
	if err != nil {
		util.Log(errors.New("数据有误"))
	}
	return key, split[1]
}
